#!/usr/bin/python

import sys
import re

tabfile = sys.argv[1]
kofile = sys.argv[2]

tf = open(tabfile, "rU")
kf = open(kofile, "rU")

tlines = tf.readlines()
klines = kf.readlines()

m1 = re.compile('(.*) \[EC:(.*)\]')
m2 = re.compile('([a-z].*)')
m3 = re.compile('(K0.*)')

ko_list = []
genes_list = []
ec_list = []
defs_list = []

for line in klines:
    l = line.split('\t')
    ko_id = l[0]
    gene_col = l[1]
    definition = l[2].rstrip()

    glist = gene_col.split(',')
    gcans = []
    g2append = ""
    for x in glist:
        y = x.strip()
        if m2.match(y):
            pattern = m2.match(y)
            gcans.append(pattern.group(1))
    if len(gcans) > 1:
        g2append = gcans[0]
    elif len(gcans) == 1:
        g2append = gcans[0]
    else:
        g2append = "-"
    genes_list.append(g2append)

    ko_list.append(ko_id)
    
    if m1.match(definition):
        pattern = m1.match(definition)
        #ec_list.append(pattern.group(2))
        defs_list.append(pattern.group(1))
        elist = pattern.group(2).split(' ')
        #ecans = []
        e2append = ""
        if len(elist) > 1:
            e2append = elist[0]
        elif len(elist) == 1:
            e2append = elist[0]
        else:
            e2append = "-"
        ec_list.append(e2append)
    else:
        ec_list.append("-")
        defs_list.append(definition)

for line in tlines:
    l = line.split('\t')
    feat = l[0]
    locus_tag = l[1]
    start = int(l[2])
    stop = int(l[3])
    frame = l[4]
    gene = l[5]
    protein_id = l[6]
    ncbi_ec_num = l[7]
    cog = l[8]
    bound_moiety = l[9]
    desc = l[10]
    ko = l[11].rstrip()
    out_str = ""
    if feat == "CDS":
        if m3.match(ko):
            pattern = m3.match(ko)
            ko2search = pattern.group(1)
            kegg_gene_name = ""
            kegg_ec_num = ""
            kegg_def_name = ""
            for index, val in enumerate(ko_list):
                if ko_list[index] == ko2search:
                    kegg_gene_name = genes_list[index]
                    kegg_ec_num = ec_list[index]
                    kegg_def_name = defs_list[index]
            out_str = locus_tag + "\t" + ko + "\t" + kegg_gene_name + "\t" + kegg_ec_num + "\t" + desc
        else:
            out_str = locus_tag + "\t" + ko + "\t" + gene + "\t" + ncbi_ec_num + "\t" + desc
    elif feat == "tRNA":
        out_str = locus_tag + "\t" + ko + "\t" + gene + "\t" + "-" + "\t" + desc
    else:
        out_str = locus_tag + "\t" + ko + "\t" + "-" + "\t" + "-" + "\t" + desc
    print out_str

tf.close()
kf.close()
